package com.dcits.components.threadpools;


import io.micrometer.core.instrument.*;
import org.springframework.boot.autoconfigure.AutoConfigureBefore;
import org.springframework.context.ApplicationContext;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.concurrent.Executor;
import java.util.concurrent.ThreadPoolExecutor;

@Component
@AutoConfigureBefore(ThreadPoolAutoConfiguration.class)  // 在其他自动配置之前加载
public class ThreadPoolMetricsBinder {
    private final MeterRegistry meterRegistry;
    private final ApplicationContext applicationContext;
    private final ExecutorRegistry executorRegistry;

    public ThreadPoolMetricsBinder(MeterRegistry meterRegistry,
                                   ApplicationContext applicationContext,
                                   ExecutorRegistry executorRegistry) {
        this.meterRegistry = meterRegistry;
        this.applicationContext = applicationContext;
        this.executorRegistry = executorRegistry;
    }

    @PostConstruct
    public void bindAllThreadPools() {
        // 绑定自定义线程池
        executorRegistry.getExecutorMap().forEach(this::bindThreadPoolMetrics);

        // 绑定Spring管理的线程池
//        applicationContext.getBeansOfType(ThreadPoolTaskExecutor.class).forEach((name, executor) -> {
//            if (!executorRegistry.getExecutorMap().containsKey(name)) {
//                bindThreadPoolMetrics(name, executor);
//            }
//        });
        executorRegistry.getExecutorMap().forEach((key, value) -> bindThreadPoolMetrics(key, value));
    }

    private void bindThreadPoolMetrics(String name, Executor executor) {
        ThreadPoolExecutor pool = executor instanceof ThreadPoolTaskExecutor
                ? ((ThreadPoolTaskExecutor) executor).getThreadPoolExecutor()
                : (ThreadPoolExecutor) executor;

        // 活动线程数
        Gauge.builder("threadpool.active.current", pool, ThreadPoolExecutor::getActiveCount)
                .tags("name", name)
                .description("当前活动线程数")
                .register(meterRegistry);

        // 已完成任务数
        FunctionCounter.builder("threadpool.completed", pool, ThreadPoolExecutor::getCompletedTaskCount)
                .tags("name", name)
                .description("已完成任务总数")
                .register(meterRegistry);

        // 当前池大小
        Gauge.builder("threadpool.pool.size", pool, ThreadPoolExecutor::getPoolSize)
                .tags("name", name)
                .description("当前线程池大小")
                .register(meterRegistry);

        // 队列大小
        Gauge.builder("threadpool.queue.size", pool, e -> e.getQueue().size())
                .tags("name", name)
                .description("当前队列中的任务数")
                .register(meterRegistry);

        // 核心线程数
        Gauge.builder("threadpool.core.size", pool, ThreadPoolExecutor::getCorePoolSize)
                .tags("name", name)
                .description("核心线程数")
                .register(meterRegistry);

        // 最大线程数
        Gauge.builder("threadpool.max.size", pool, ThreadPoolExecutor::getMaximumPoolSize)
                .tags("name", name)
                .description("最大线程数")
                .register(meterRegistry);
    }
}